#自动切图

import glob
import os
import shutil
import xml.etree.cElementTree as et
import cv2
import fire
import numpy as np
import sys
import re
from enum import Enum
import time
import logging
from scipy import signal
from scipy.signal import argrelmin, argrelmax
from scipy.spatial import distance as dist
import matplotlib.pyplot as plt
import xml.dom.minidom as minidom
import pandas as pd



def get_cut_line(img, cols, rows, xthre, ythre):
    '''
    寻找切线。
    '''
    t = img[:,:,0]
    hh, ww = t.shape[:2]
    xstep = ww // cols
    ystep = hh // rows


    horizontal = np.sum(t, axis = 0)
    vertical = np.sum(t, axis = 1)


    # 电池片水平栅格线有可能干扰到电池片的水平切割，特别处理下
    conv_width_half = hh // (rows * 90)
    vertical = np.convolve(vertical, np.ones(conv_width_half*2 +1).astype('int'))[conv_width_half:-1*conv_width_half]
    # 计算梯度
    hg = np.gradient(horizontal)
    vg = np.gradient(vertical)


    col_pos = []
    row_pos = []
    # 采用梯度过滤

    col_lines = [0]
    xstart = xstep
    for x in range(cols):
        xend = np.argmax(hg[xstart: xstart + xstep//2]) + xstart
        xstart = np.argmin(hg[xstart: xstart + xstep//2]) + xstart
        if x == cols//2 - 1:
            # 汇流条处有两条切线
            col_lines.append(int(xend))
            col_lines.append(int(xstart))
            xstart += xstep
        else:
            if xstart - xend < xthre:
                col_pos.append(len(col_lines))
            col_lines.append(int(np.mean([xend, xstart])))
            xstart += xstep*3//4


    row_lines = [0]
    ystart = ystep * 3 // 4
    for y in range(rows):
        yend = np.argmax(vg[ystart: ystart + ystep//2]) + ystart
        ystart = np.argmin(vg[ystart: ystart + ystep//2]) + ystart
        if y == rows-1:
            row_lines.append(int(yend+1))
        else:
            if ystart - yend < ythre:
                row_pos.append(len(row_lines))
            row_lines.append(int(np.mean([yend, ystart])))
        ystart += ystep*3//4

    hsep = (col_lines[2] - col_lines[1])
    col_lines[0] = max(col_lines[1] - hsep, 0)
    col_lines[-1] = min(col_lines[-2] + hsep, ww)

    vsep = (row_lines[2] - row_lines[1])
    # 如果最下方的切线不准，尝试使用最上方切线
    if row_lines[1] - vsep < 0 :
        vsep = (row_lines[1] - row_lines[0])
        row_lines[0] = max(row_lines[1] - vsep, 0)
        row_lines[-1] = min(row_lines[-2] + vsep, hh)
    else:
        row_lines[0] = max(row_lines[1] - vsep, 0)
        row_lines[-1] = min(row_lines[-2] + vsep, hh)


    return col_lines, row_lines, col_pos, row_pos



def get_cut_img(img, img_name, col_lines, row_lines, col_pos, row_pos, resever):
    if resever:
        rn = ['F', 'E', 'D', 'C', 'B', 'A']
    else:
        rn = ['A', 'B', 'C', 'D', 'E', 'F']
    chuancut = {}
    piancut = {}

    # 切串
    for c in range(len(col_lines)-1): 
        # 汇流条切线有两条，列总切线数量为列数加2，遍历到倒数第二根切线
        if c == (len(col_lines)-2)//2:
            continue
        for r in range(len(row_lines)-2):
            # 行总切线数量为行数加1，遍历到倒数第三根切线
            img_cutchuan = cv2.resize(img[row_lines[r]:row_lines[r+2], col_lines[c]:col_lines[c+1], :], (375, 1500))
            if c > (len(col_lines)-2)//2:
                img_chuanname = ''.join([img_name, '_', str(rn[r]), str(c), str(rn[r+1]), str(c)])
            else:
                img_chuanname = ''.join([img_name, '_', str(rn[r]), str(c+1), str(rn[r+1]), str(c+1)])
            chuancut[img_chuanname] = img_cutchuan

    # 切片
    for r in range(len(row_lines)-1):
        # 行总切线数量为行数加1，遍历到倒数第二根切线
        for c in range(len(col_lines)-2): 
            # 汇流条切线有两条，列总切线数量为列数加2，遍历到倒数第三根切线
            if c == (len(col_lines)-2)//2 -1:
                continue
            if c == (len(col_lines)-2)//2:
                continue
            img_cutpian = cv2.resize(np.rot90(img[row_lines[r]:row_lines[r+1], col_lines[c]:col_lines[c+2], :]), (750, 750))
            if c > (len(col_lines)-2)//2:
                img_pianname = ''.join([img_name, '_', str(rn[r]), str(c), str(rn[r]), str(c+1)])
            else:
                img_pianname = ''.join([img_name, '_', str(rn[r]), str(c+1), str(rn[r]), str(c+2)])
            piancut[img_pianname] = img_cutpian

    return chuancut, piancut



def cut_img_save(pian_save_path, chuan_save_path, chuancut, piancut):
    for chuan_img in chuancut:
        save_file = os.path.join(chuan_save_path, ''.join([chuan_img, '.jpg']))
        cv2.imwrite(save_file, chuancut[chuan_img])

    for pian_img in piancut:
        save_file = os.path.join(pian_save_path, ''.join([pian_img, '.jpg']))
        cv2.imwrite(save_file, piancut[pian_img])



def pos_mat(save_path, pos_file):
    pos_dic = {}
    img_li = []
    with open(pos_file) as f:
        line = f.readline()
        while line:
            line_li = line.split(' ')
            pos_dic[line_li[0]] = line_li[-1]
            line = f.readline()
    for img_name in pos_dic:
        for img_file in glob.glob(os.path.join(save_path, ''.join([img_name, '*']))):
            img_pos = os.path.basename(img_file).split('.jpg').split('_')[-1]
            if pos_dic[img_name] in img_pos:
                img_li.append(img_file)
    return img_li



def check_dir(dir_path):
    '''
    检测路径，如果有路径提示是否删除，无路径新建路径
    '''
    if os.path.exists(dir_path):
        print('!!!dir {} exists'.format(dir_path))
        comfirm = input("!!!remove dir, yes or mix? [y or m]: ")
        if comfirm == 'y':
            print('!!!{} have been removed!!!'.format(dir_path))
            shutil.rmtree(dir_path)
            os.makedirs(dir_path)
        elif comfirm != 'm':
            print('did not remove!')
            sys.exit(0)
    else:
        os.makedirs(dir_path)



def main(img_path, save_path, pos_file, cols = 24, rows = 6, xthre = 6, ythre = 6, resever = True, sw = False):
    chuan_save_path = os.path.join(save_path, 'chuan')
    pian_save_path = os.path.join(save_path, 'pian')
    check_dir(chuan_save_path)
    check_dir(pian_save_path)
    for img_file in glob.glob(os.path.join(img_path, '*.jpg')):
        img_name = os.path.basename(img_file).split('.jpg')[0]
        print('<===============%===============>')
        print('{} is processing'.format(img_name))
        img = cv2.imdecode(np.fromfile(img_file, dtype=np.uint8), 1)
        col_lines, row_lines, col_pos, row_pos = get_cut_line(img, cols, rows, xthre, ythre)
        chuancut, piancut = get_cut_img(img, img_name, col_lines, row_lines, col_pos, row_pos, resever)
        cut_img_save(pian_save_path, chuan_save_path, chuancut, piancut)
        print('pian:{}, chuan:{}, cutted'.format(len(piancut.keys()), len(chuancut.keys())))

    if sw:
        have_path = os.path.join(save_path, 'all_have')
        chuan_img_li = pos_mat(chuan_save_path, pos_file)
        pian_img_li = pos_mat(pian_save_path, pos_file)
        for img_file in chuan_img_li:
            shutil.copy(img_file, have_path)
        for img_file in pian_img_li:
            shutil.copy(img_file, have_path)


if __name__ == '__main__':
    fire.Fire(main)


